import numpy

def set_up_datasets(args):
    if args.dataset == 'miniimagenet':
        args.num_class = 64
        if args.deepemd == 'fcn':
            from Models.dataloader.miniimagenet.fcn.mini_imagenet import MiniImageNet as Dataset
        elif args.deepemd == 'sampling':
            from Models.dataloader.miniimagenet.sampling.mini_imagenet import MiniImageNet as Dataset
        elif args.deepemd == 'grid':
            from Models.dataloader.miniimagenet.grid.mini_imagenet import MiniImageNet as Dataset
    elif args.dataset == 'cub':
        args.num_class = 64
        if args.deepemd == 'fcn':
            from Models.dataloader.cub.fcn.cub import CUB as Dataset
        elif args.deepemd == 'sampling':
            from Models.dataloader.cub.sampling.cub import CUB as Dataset
        elif args.deepemd == 'grid':
            from Models.dataloader.cub.grid.cub import CUB as Dataset
    elif args.dataset == 'fc100':
        args.num_class = 60
        if args.deepemd == 'fcn':
            from Models.dataloader.fc100.fcn.fc100 import DatasetLoader as Dataset
        elif args.deepemd == 'sampling':
            from Models.dataloader.fc100.sampling.fc100 import DatasetLoader as Dataset
        elif args.deepemd == 'grid':
            from Models.dataloader.fc100.grid.fc100 import DatasetLoader as Dataset
    elif args.dataset == 'tieredimagenet':
        args.num_class = 351
        if args.deepemd == 'fcn':
            from Models.dataloader.tieredimagenet.fcn.tiered_imagenet import tieredImageNet as Dataset
        elif args.deepemd == 'sampling':
            from Models.dataloader.tieredimagenet.sampling.tiered_imagenet import tieredImageNet as Dataset
        elif args.deepemd == 'grid':
            from Models.dataloader.tieredimagenet.grid.tiered_imagenet import tieredImageNet as Dataset
    elif args.dataset == 'cifar_fs':
        args.num_class = 64
        if args.deepemd == 'fcn':
            from Models.dataloader.cifar_fs.fcn.cifar_fs import DatasetLoader as Dataset
        elif args.deepemd == 'sampling':
            from Models.dataloader.cifar_fs.sampling.cifar_fs import DatasetLoader as Dataset
        elif args.deepemd == 'grid':
            from Models.dataloader.cifar_fs.gird.cifar_fs import DatasetLoader as Dataset
    elif args.dataset == 'recognition36':
        args.num_class = 20
        if args.deepemd == 'fcn':
            from Models.dataloader.recognition36.fcn.recognition_36 import Recognition36 as Dataset
        elif args.deepemd == 'sampling':
            from Models.dataloader.recognition36.sampling.recognition_36 import Recognition36 as Dataset
        elif args.deepemd == 'grid':
            from Models.dataloader.recognition36.grid.recognition_36 import Recognition36 as Dataset
    elif args.dataset == 'recognition36_crop':
        args.num_class = 20
        if args.deepemd == 'fcn':
            from Models.dataloader.recognition36_crop.fcn.recognition36_crop import recognition36Crop as Dataset
        elif args.deepemd == 'sampling':
            from Models.dataloader.recognition36_crop.sampling.recognition36_crop import recognition36Crop as Dataset
        elif args.deepemd == 'grid':
            from Models.dataloader.recognition36_crop.grid.recognition36_crop import recognition36Crop as Dataset
    elif args.dataset == 'cars':
        args.num_class = 100
        if args.deepemd == 'fcn':
            from Models.dataloader.cars.fcn.cars import cars as Dataset
        elif args.deepemd == 'sampling':
            from Models.dataloader.cars.sampling.cars import cars as Dataset
        elif args.deepemd == 'grid':
            from Models.dataloader.cars.grid.cars import cars as Dataset
    elif args.dataset == 'VHR-10':
        args.num_class = 10
        if args.deepemd == 'fcn':
            from Models.dataloader.VHR_10.fcn.vhr_10 import VHR_10 as Dataset
        elif args.deepemd == 'sampling':
            pass
        elif args.deepemd == 'grid':
            pass
    else:
        raise ValueError('Unkown Dataset')
    return Dataset

def compute_mean_std(dataset):
    """compute the mean and std of cifar100 dataset
    Args:
        cifar100_training_dataset or cifar100_test_dataset
        witch derived from class torch.utils.data

    Returns:
        a tuple contains mean, std value of entire dataset
    """
    # RGB 三通道
    if len(numpy.array(dataset[0][0]).shape) == 3:
        data = numpy.stack([numpy.array(dataset[i][0]) for i in range(len(dataset))], axis=-1)
        data_r = data[:,:,0,:]
        data_g = data[:,:,1,:]
        data_b = data[:,:,2,:]
        mean = numpy.mean(data_r), numpy.mean(data_g), numpy.mean(data_b)
        std = numpy.std(data_r), numpy.std(data_g), numpy.std(data_b)
    # 单通道灰度图
    else:
        data = numpy.dstack([numpy.array(dataset[i][0]) for i in range(len(dataset))])
        mean = numpy.mean(data)
        std = numpy.std(data)
    return mean, std


if __name__ == "__main__":
    from recognition36.fcn.recognition_36 import Recognition36
    class Args:
        pass
    args = Args()
    args.image_size = 84
    args.data_dir = "datasets/"
    dataset = Recognition36(setname="test", args=args)
    compute_mean_std(dataset)
    print()